#!/usr/bin/env python

import rospy
from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError
import cv2
import select
import sys
import tty
import termios
import rospkg

def getKey():
    old_settings = termios.tcgetattr(sys.stdin)
    try:
        tty.setcbreak(sys.stdin.fileno())
        if select.select([sys.stdin], [], [], 0)[0]:
            key = sys.stdin.read(1)
        else:
            key = None
    finally:
        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings)
    return key

class CameraSubscriber:
    def __init__(self):
        self.bridge = CvBridge()
        self.image_sub = rospy.Subscriber("/kinect2/qhd/image_color_rect", Image, self.callback)
        self.cv_image = None

    def callback(self, data):
        try:
            self.cv_image = self.bridge.imgmsg_to_cv2(data, "bgr8")
        except CvBridgeError as e:
            print(e)
            return

        # 在这里处理图像，例如显示图像
        cv2.imshow("Camera Image", self.cv_image)
        cv2.waitKey(1)

if __name__ == '__main__':
    rospy.init_node('camera_subscriber', anonymous=True)
    camera_subscriber = CameraSubscriber()
    # 获取rospack实例
    rospack = rospkg.RosPack()
    # 找到cv_pkg包的路径
    cv_pkg_path = rospack.get_path('cv_pkg')
    ImgNum = 0
    while not rospy.is_shutdown():
        key = getKey()
        if key == 'j' and camera_subscriber.cv_image is not None:  # 按j保存一张图片
            ImgNum += 1
            u = str(ImgNum)
            filename = cv_pkg_path + '/data/img' + u + '.jpg'
            cv2.imwrite(filename, camera_subscriber.cv_image)
            print('写入：', filename)
        elif key == 'q':
            cv2.destroyAllWindows()
            break


    try:
        rospy.spin()
    except KeyboardInterrupt:
        print("Shutting down")
